#include "enonbond.h"


namespace kuai { namespace cuda {

	__global__ void setuplist(const int input[], int n, int rank, int skip,
		int output[], int counter[])
	{
		int tid = THREAD_ID;
		if (tid < n) {
			int i0 = MUL(tid, skip);			// start index of input
			int o0 = MUL(tid, skip+rank);	// start index of output
			for (int i = 0; i < rank; ++i) {
				int atomI = output[o0+i] = input[i0+i];
				output[o0+skip+i] = atomicAdd(counter+atomI, 1);
			}
			for (int i = rank; i < skip; ++i) {
				output[o0+i] = input[i0+i];
			}
		}
	}
	
	__device__ inline void 
		getIndexAndR(
			int tid,
			const int list[], int skipList, 
			const RealNumber coord[], int skipCoord,
			int index[], RealNumber r[]
		) 
	{
		int l0 = MUL(tid, skipList);
		index[0] = list[l0+0];		// atom1
		index[1] = list[l0+1];		// atom2
		index[2] = list[l0+2];		// ki
		index[3] = list[l0+2]*3;	// result1
		index[4] = list[l0+2]*3;	// result2
		int coord1 = MUL(index[0], skipCoord);
		int coord2 = MUL(index[1], skipCoord);
		
		r[0] = coord[coord1+0] - coord[coord2+0];
		r[1] = coord[coord1+1] - coord[coord2+1];
		r[2] = coord[coord1+2] - coord[coord2+2];
	};
	
	__device__ inline void 
		updateForces(
			const int index[],
			RealNumber r[], const RealNumber dedr,
			RealNumber forces[]
		) 
	{
		r[0] *= dedr;  r[1] *= dedr; r[2] *= dedr; 
		forces[index[3]+0] = r[0]; forces[index[4]+0] = -r[0];
		forces[index[3]+1] = r[1]; forces[index[4]+1] = -r[1];
		forces[index[3]+2] = r[2]; forces[index[4]+2] = -r[2];
	};
	
	__global__ void e_n12_6(const int list[], int n, int skipList, 
		const RealNumber coord[], int skipCoord, const RealNumber k[], 
		RealNumber energies[])
	{
		int tid = THREAD_ID;
		if (tid < n) {
			int index[5];
			RealNumber r[3];
			getIndexAndR(tid, list, skipList, coord, skipCoord, index, r);
			RealNumber r0 = k[index[2]];
			RealNumber rr6 = r0*r0/(r[0]*r[0]+r[1]*r[1]+r[2]*r[2]);
			rr6 = rr6 * rr6 * rr6;
			RealNumber e0 = k[index[2]+1];
			energies[tid] = e0 * rr6 * (rr6-2);
		}
	}
	
	__global__ void f_n12_6(const int list[], int n, int skipList, 
		const RealNumber coord[], int skipCoord, const RealNumber k[], 
		RealNumber energies[], RealNumber forces[])
	{
		int tid = THREAD_ID;
		if (tid < n) {
			int index[3];
			RealNumber r[3];
			getIndexAndR(tid, list, skipList, coord, skipCoord, index, r);
			RealNumber rr2 = 1/(r[0]*r[0]+r[1]*r[1]+r[2]*r[2]);
			RealNumber r0 = k[index[2]];
			RealNumber rr6 = r0*r0*rr2;
			rr6 = rr6 * rr6 * rr6;
			RealNumber e0 = k[index[2]+1];
			energies[tid] = e0 * rr6 * (rr6-2);
			RealNumber dedr = 12 * e0 * rr6 * (rr6 - 1) * rr2;
			
			updateForces(index, r, dedr, forces);			
		}
	}
	
	__global__ void e_coulomb(const int list[], int n, int skipList, 
		const RealNumber coord[], int skipCoord, 
		const RealNumber charges[], RealNumber coulombConstant,
		RealNumber energies[])
	{
		int tid = THREAD_ID;
		int l0 = MUL(tid, skipList);
		int atom1 = list[l0+0];
		int atom2 = list[l0+1];
		
		int coord1 = MUL(atom1, skipCoord);
		int coord2 = MUL(atom2, skipCoord);
		
		RealNumber x[] = {
			coord[coord1+0] - coord[coord2+0],
			coord[coord1+1] - coord[coord2+1],
			coord[coord1+2] - coord[coord2+2]};
		RealNumber rr = 1/sqrt(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]);
		energies[tid] = coulombConstant * charges[atom1] * charges[atom2] * rr;
	}
	
	__global__ void f_coulomb(const int list[], int n, int skipList, 
		const RealNumber coord[], int skipCoord, 
		const RealNumber charges[], RealNumber coulombConstant,
		RealNumber energies[], RealNumber forces[])
	{
		int tid = THREAD_ID;
		int l0 = MUL(tid, skipList);
		int atom1 = list[l0+0];
		int atom2 = list[l0+1];
		int result1 = list[l0+3]*3;
		int result2 = list[l0+4]*3;
		
		int coord1 = MUL(atom1, skipCoord);
		int coord2 = MUL(atom2, skipCoord);
		
		RealNumber x[] = {
			coord[coord1+0] - coord[coord2+0],
			coord[coord1+1] - coord[coord2+1],
			coord[coord1+2] - coord[coord2+2]};
		RealNumber rr = 1/sqrt(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]);
		RealNumber rr2 = rr*rr;
		RealNumber e = coulombConstant * charges[atom1] * charges[atom2] * rr;
		energies[tid] = e;
		x[0] *= e * rr2; x[1] *= e * rr2; x[2] *= e * rr2;
		forces[result1+0] = x[0]; forces[result2+0] = -x[0];
		forces[result1+1] = x[1]; forces[result2+1] = -x[1];
		forces[result1+2] = x[2]; forces[result2+2] = -x[2];
	}
	
}}
